package assignment04;

import gnu.trove.iterator.TIntIterator;
import gnu.trove.list.TIntList;
import gnu.trove.list.array.TIntArrayList;

import java.util.List;

import nlp.util.Indexer;

import assignment04.WordAlignmentTester.Alignment;
import assignment04.WordAlignmentTester.SentencePair;
import assignment04.WordAlignmentTester.WordAligner;

public class IBMTwoAligner implements WordAligner {
	private final static int NULL_INDEX = -1;
	private final static double PNULL = .2;
	private final static double PREST = .8;
	
	public static class IBMTwoFactory implements WordAlignerFactory {
		IntCounter frenchCounter;
		IntCounter englishCounter;
		IntCounterMap cMap;
		IntCounterMap tMap;
		IntCounterMap another_cMap;
		IntCounterMap aMap;
		
		Indexer<String> englishIndexer;
		Indexer<String> frenchIndexer;
		Indexer<TIntList> positionIndexer;
		
		public WordAligner trainAligner(List<SentencePair> trainingSentencePairs) {
			buildIndexerAndCounter(trainingSentencePairs);
			///**
			for (int i = 0; i < 5; i++)
			{
				System.err.println("model one training loop " + i);
				modelOneTrainTStep(trainingSentencePairs);
				modelOneTrainMStep();
			}
			//*/
			
			for (int i = 0; i < 5; i++) {
				System.err.println("model two training loop " + i);
				modelTwoTrainTStep(trainingSentencePairs);
				modelTwoTrainMStep();
			}
			
			return new IBMTwoAligner(frenchCounter, englishCounter, tMap, aMap, englishIndexer, frenchIndexer, positionIndexer);
		}
		
		private void modelTwoTrainTStep(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;
				modelTwoInnerTrain(eList, fList);
			}
		}
		
		private void modelTwoInnerTrain(List<String> eList, List<String> fList) {
			int numFrench = fList.size();
			int numEnglish = eList.size();
			
			for (int frenchPosition = 0; frenchPosition < numFrench; frenchPosition++) {
				double denominator = 0.0;
				int frenchIndex = frenchIndexer.indexOf(fList.get(frenchPosition));
				int positionIndex = positionIndexer.indexOf(new TIntArrayList(new int[]{numFrench, numEnglish, frenchPosition}));
				for (int englishPosition = -1/*TODO*/; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = -1;
					if (0 <= englishPosition) {
						englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					}
					
					denominator += aMap.getCount(positionIndex, englishPosition)
								* tMap.getCount(englishIndex, frenchIndex);
				}
				
				for (int englishPosition = -1/*TODO*/; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = -1;
					if (0 <= englishPosition) {
						englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					}
					double p = aMap.getCount(positionIndex, englishPosition)
							* tMap.getCount(englishIndex, frenchIndex)
							/ denominator;
//					System.err.println("p is " + aMap.getCount(positionIndex, englishPosition) + " and " + tMap.getCount(englishIndex, frenchIndex));
					cMap.incrementCount(englishIndex, frenchIndex, p);
					another_cMap.incrementCount(positionIndex, englishPosition, p);
				}
			}
		}
		
		private void modelTwoTrainMStep() {
			tMap = IntCounters.conditionalNormalize(cMap);
			cMap.reset();
			aMap = IntCounters.conditionalNormalize(another_cMap);
			another_cMap.reset();
		}

		
		private void buildIndexerAndCounter(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;

				for (String eWord : eList) {
					englishCounter.incrementCount(englishIndexer.addAndGetIndex(eWord), 1.0);
				}
				
				for (String fWord : fList) {
					int fIndex = frenchIndexer.addAndGetIndex(fWord);
					frenchCounter.incrementCount(fIndex, 1.0);
					cMap.setCount(NULL_INDEX, fIndex, 0.0);
				}
				
				for (String english : eList) {
					int eIndex = englishIndexer.indexOf(english);
					for (String french : fList) {
						int fIndex = frenchIndexer.indexOf(french);
						cMap.setCount(eIndex, fIndex, 0.0);
					}
				}
				
				int eLength = eList.size();
				int fLength = fList.size();
				for (int frenchPosition = 0; frenchPosition < fLength; frenchPosition++) {
					TIntList intList = new TIntArrayList(new int[]{fLength, eLength, frenchPosition});
					int positionIndex = positionIndexer.addAndGetIndex(intList);
					for (int englishPosition = -1; englishPosition < eLength; englishPosition++) {
						another_cMap.setCount(positionIndex, englishPosition, 0.0);
						aMap.setCount(positionIndex, englishPosition, 1.0);
					}
				}
			}
			
			//normalize
			englishCounter.normalize();
			frenchCounter.normalize();
			
			System.out.println("there are " + englishCounter.size() + " english word types");
			System.out.println("there are " + frenchCounter.size() + " french word types");
			
			for (TIntIterator outerIt = cMap.keySet().iterator(); outerIt.hasNext();) {
				int key = outerIt.next();
				IntCounter counter = cMap.getCounter(key);
				for (TIntIterator innerIt = counter.keySet().iterator(); innerIt.hasNext();) {
					int value = innerIt.next();
					tMap.setCount(key, value, frenchCounter.getCount(value));
				}
			}
		}
		
		
		private void modelOneTrainTStep(List<SentencePair> trainingSentencePairs) {
			for (SentencePair pair : trainingSentencePairs) {
				List<String> eList = pair.englishWords;
				List<String> fList = pair.frenchWords;
				modelOneInnerTrain(eList, fList);
			}
		}
		
		private void modelOneTrainMStep() {
			tMap = IntCounters.conditionalNormalize(cMap);
			cMap.reset();
		}
		
		private void modelOneInnerTrain(List<String> eList, List<String> fList) {
			int numFrench = fList.size();
			int numEnglish = eList.size();
			
			for (int frenchPosition = 0; frenchPosition < numFrench; frenchPosition++) {
				int frenchIndex = frenchIndexer.indexOf(fList.get(frenchPosition));
				double tNull = tMap.getCount(NULL_INDEX, frenchIndex);
				double tSum = 0.0;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					tSum += tMap.getCount(englishIndex, frenchIndex);
				}
				double denorminator = (PNULL * tNull) + PREST / numEnglish * tSum;
				for (int englishPosition = 0; englishPosition < numEnglish; englishPosition++) {
					int englishIndex = englishIndexer.indexOf(eList.get(englishPosition));
					double pNormal = (PREST / numEnglish * tMap.getCount(englishIndex, frenchIndex))
							/ denorminator;
							//((PNULL * tNull) + PREST / numEnglish * tSum);
					cMap.incrementCount(englishIndex, frenchIndex, pNormal);
				}
				double pSpecial = (PNULL * tNull)
								/ denorminator;
								//(PNULL * tNull + PREST / numEnglish * tSum);
				cMap.incrementCount(NULL_INDEX, frenchIndex, pSpecial);
			}
		}
		
		public IBMTwoFactory() {
			frenchCounter = new IntCounter();
			englishCounter = new IntCounter();
			cMap = new IntCounterMap();
			tMap = new IntCounterMap();
			another_cMap = new IntCounterMap();
			aMap = new IntCounterMap();
			
			englishIndexer = new Indexer<String>();
			frenchIndexer = new Indexer<String>();
			positionIndexer = new Indexer<TIntList>();
		}
	}
	
	
	public Alignment alignSentencePair(SentencePair sentencePair) {
		
		Alignment alignment = new Alignment();
		int numFrenchWords = sentencePair.getFrenchWords().size();
		int numEnglishWords = sentencePair.getEnglishWords().size();
		for (int frenchPosition = 0; frenchPosition < numFrenchWords; frenchPosition++) {
			int maxPosition = -1;
			double maxScore = Double.NEGATIVE_INFINITY;
			String frenchWord = sentencePair.frenchWords.get(frenchPosition);
			int frenchIndex = frenchIndexer.indexOf(frenchWord);
			int positionIndex = positionIndexer.indexOf(new TIntArrayList(new int[]{numFrenchWords, numEnglishWords, frenchPosition}));
			
			for (int englishPosition = -1; // TODO don't forget to set to -1
			englishPosition < numEnglishWords; englishPosition++) {
				int englishIndex = -1;
				if (englishPosition >= 0) {
					String englishWord = sentencePair.englishWords
							.get(englishPosition);
					englishIndex = englishIndexer.indexOf(englishWord);
				}
				double score = aMap.getCount(positionIndex, englishPosition)
							* tMap.getCount(englishIndex, frenchIndex);
//				System.out.println("score is " + aMap.getCount(positionIndex, englishPosition)
//						+ " multiply by " +  tMap.getCount(englishIndex, frenchIndex));
				if (score > maxScore) {
					maxScore = score;
					maxPosition = englishPosition;
				}
			}
			alignment.addAlignment(maxPosition, frenchPosition, true);
		}
		return alignment;
	}
	
	public IBMTwoAligner(IntCounter frenchCounter, IntCounter englishCounter,
			IntCounterMap tMap, IntCounterMap aMap, Indexer<String> englishIndexer,
			Indexer<String> frenchIndexer, Indexer<TIntList> positionIndexer) {
		this.frenchCounter = frenchCounter;
		this.englishCounter = englishCounter;
		this.tMap = tMap;
		this.aMap = aMap;
		this.englishIndexer = englishIndexer;
		this.frenchIndexer = frenchIndexer;
		this.positionIndexer = positionIndexer;

	}
	
	IntCounter frenchCounter;
	IntCounter englishCounter;
	IntCounterMap tMap;
	IntCounterMap aMap;
	Indexer<String> englishIndexer;
	Indexer<String> frenchIndexer;
	Indexer<TIntList> positionIndexer;


}
